from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.messages import  TextMessage
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
#from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.models.ollama import OllamaChatCompletionClient
import asyncio
# Assuming your Ollama server is running locally on port 11434.
model_client = OllamaChatCompletionClient(
        #model="hhao/qwen2.5-coder-tools:latest",
        model="qwen2.5-coder:latest",
        #model="GandalfBaum/deepseek_r1-claude3.7:latest",
        #model = "GandalfBaum/llama3.2-claude3.7:latest",
        host="http://192.168.99.142:11434", 
        model_info={
        "vision": False,
        "function_calling": True,
        "json_output": True,
        "family": "unknow",
        "structed_output":True,
        },
    )
# Define a tool that searches the web for information.
async def web_search(query: str) -> str:
    """Find information on the web """
    return " AutoGen is a programming framework for building multi-agent applications. "

agent = AssistantAgent(
    name="assistant",
    model_client=model_client,
    tools=[web_search],
    system_message="You are a helpful assistant. You can call tools to help user and  solve tasks.",
    reflect_on_tool_use= True, 
)
"""
async def assistant_run() -> None:
    response = await agent.on_messages(
        [TextMessage(content="Find information on AutoGen", source="user")],
        cancellation_token=CancellationToken(),
    )
    print(response.inner_messages)
    print(response.chat_message)


asyncio.run(assistant_run())

"""
async def assistant_run_stream() -> None:
    # Option 1: read each message from the stream (as shown in the previous example).
    # async for message in agent.on_messages_stream(
    #     [TextMessage(content="Find information on AutoGen", source="user")],
    #     cancellation_token=CancellationToken(),
    # ):
    #     print(message)

    # Option 2: use Console to print all messages as they appear.
    await Console(
        agent.on_messages_stream(
            [TextMessage(content="Find information on AutoGen", source="user")],
            cancellation_token=CancellationToken(),
        ),
        output_stats=True,  # Enable stats printing.
    )


asyncio.run(assistant_run_stream()) 

